#include <algorithm>

#include "scene.hpp"

void Scene::box(
    std::vector<FaceAndCenter>::const_iterator begin,
    std::vector<FaceAndCenter>::const_iterator end,
    Eigen::Vector3f & min,
    Eigen::Vector3f & max) const
{
    min = Eigen::Vector3f::Constant(INFINITY);
    max = Eigen::Vector3f::Constant(-INFINITY);

    for (auto it = begin; it != end; ++it)
    {
        Eigen::Matrix3f Ps;
        objects[it->face.object_idx].face_vertices(it->face.face_idx, Ps);

        min = min.cwiseMin(Ps.rowwise().minCoeff());
        max = max.cwiseMax(Ps.rowwise().maxCoeff());
    }
}

void Scene::partition(
    int node,
    std::vector<FaceAndCenter>::iterator begin,
    std::vector<FaceAndCenter>::iterator end,
    int level)
{
    box(begin, end, tree[node].min, tree[node].max);

    if (level == 0)
    {
        auto & leaf_faces = faces[node + 1 - faces.size()];
        for (auto it = begin; it != end; ++it)
            leaf_faces.push_back(it->face);
        return;
    }

    int axis;
    (tree[node].max - tree[node].min).maxCoeff(&axis);

    std::sort(
        begin,
        end,
        [axis](const FaceAndCenter & a, const FaceAndCenter & b)
        {
            return a.center(axis) < b.center(axis);
        }
    );

    int half_len = (end - begin) / 2;

    partition(node * 2 + 1, begin, begin + half_len, level - 1);
    partition(node * 2 + 2, begin + half_len, end, level - 1);
}

void Scene::intersect(
    int node,
    const Eigen::Vector3f & ray_o,
    const Eigen::Vector3f & ray_d,
    float & res_dist,
    Face & res_face) const
{
    Eigen::Vector3f t_a = (tree[node].min - ray_o).array() / ray_d.array();
    Eigen::Vector3f t_b = (tree[node].max - ray_o).array() / ray_d.array();

    Eigen::Vector3f t_min = t_a.cwiseMin(t_b);
    Eigen::Vector3f t_max = t_a.cwiseMax(t_b);

    float t_enter = t_min.maxCoeff();
    float t_exit = t_max.minCoeff();

    if (t_enter < t_exit && t_exit > 0) // intersects with box
        if (node + 1 >= faces.size()) // leaf node
        {
            for (auto & face : faces[node + 1 - faces.size()])
            {
                Eigen::Matrix3f Ps;
                objects[face.object_idx].face_vertices(face.face_idx, Ps);

                Eigen::Vector3f E1 = Ps.col(1) - Ps.col(0);
                Eigen::Vector3f E2 = Ps.col(2) - Ps.col(0);
                Eigen::Vector3f S = ray_o - Ps.col(0);
                Eigen::Vector3f S1 = ray_d.cross(E2);
                Eigen::Vector3f S2 = S.cross(E1);
                float k = 1.0f / S1.dot(E1);

                float t = k * S2.dot(E2);
                float b1 = k * S1.dot(S);
                float b2 = k * S2.dot(ray_d);

                if (t > 1e-5 && t < res_dist && b1 > 0 && b2 > 0 && b1 + b2 < 1)
                {
                    res_dist = t;
                    res_face = face;
                }
            }
        }
        else
        {
            intersect(node * 2 + 1, ray_o, ray_d, res_dist, res_face);
            intersect(node * 2 + 2, ray_o, ray_d, res_dist, res_face);
        }
}

Scene::Scene(std::vector<Object> && objects_, SkyFn sky_):
    objects(objects_), sky(sky_)
{
    std::vector<FaceAndCenter> centers;

    for (int i = 0; i < objects.size(); i++)
        for (int j = 0; j < objects[i].faces.rows(); j++)
        {
            Eigen::Matrix3f Ps;
            objects[i].face_vertices(j, Ps);
            centers.push_back({{i, j}, Ps.rowwise().mean()});
        }

    int level = 0, leaves = 1;
    while (leaves < centers.size() / 5)
    {
        level++;
        leaves *= 2;
    }

    tree.resize(leaves * 2 - 1);
    faces.resize(leaves);

    partition(0, centers.begin(), centers.end(), level);
}

void Scene::intersect(
    const Eigen::Vector3f & ray_o,
    const Eigen::Vector3f & ray_d,
    float & dist,
    int & object_idx,
    int & face_idx,
    Eigen::Vector3f & sky_color) const
{
    dist = INFINITY;
    Face face{-1, -1};

    intersect(0, ray_o, ray_d, dist, face);

    object_idx = face.object_idx;
    face_idx = face.face_idx;

    if (object_idx < 0)
        sky_color = sky(ray_d);
}
